import os
import pandas as pd
from django.core.management.base import BaseCommand
from django.contrib.auth.models import User
from movies.models import Genre, Movie, Rating
from django.db import transaction
import re


class Command(BaseCommand):
    help = '加载MovieLens数据集到数据库'
    
    def add_arguments(self, parser):
        parser.add_argument(
            '--data-path',
            type=str,
            default=r'D:\ml-latest-small\ml-latest-small',
            help='MovieLens数据集路径'
        )
    
    def handle(self, *args, **options):
        data_path = options['data_path']
        
        # 检查数据文件是否存在
        movies_file = os.path.join(data_path, 'movies.csv')
        ratings_file = os.path.join(data_path, 'ratings.csv')
        
        if not os.path.exists(movies_file):
            self.stdout.write(
                self.style.ERROR(f'电影文件不存在: {movies_file}')
            )
            return
            
        if not os.path.exists(ratings_file):
            self.stdout.write(
                self.style.ERROR(f'评分文件不存在: {ratings_file}')
            )
            return
        
        self.stdout.write('开始加载MovieLens数据...')
        
        # 加载电影数据
        self.load_movies(movies_file)
        
        # 加载评分数据
        self.load_ratings(ratings_file)
        
        self.stdout.write(
            self.style.SUCCESS('MovieLens数据加载完成!')
        )
    
    def load_movies(self, movies_file):
        """加载电影数据"""
        self.stdout.write('正在加载电影数据...')
        
        # 读取电影数据
        movies_df = pd.read_csv(movies_file)
        
        with transaction.atomic():
            for _, row in movies_df.iterrows():
                movie_id = row['movieId']
                title = row['title']
                genres_str = row['genres']
                
                # 提取年份
                year_match = re.search(r'\((\d{4})\)$', title)
                year = int(year_match.group(1)) if year_match else None
                
                # 创建或获取电影
                movie, created = Movie.objects.get_or_create(
                    movie_id=movie_id,
                    defaults={
                        'title': title,
                        'year': year
                    }
                )
                
                if created:
                    # 处理类型
                    if genres_str and genres_str != '(no genres listed)':
                        genre_names = genres_str.split('|')
                        for genre_name in genre_names:
                            genre, _ = Genre.objects.get_or_create(name=genre_name)
                            movie.genres.add(genre)
        
        self.stdout.write(
            self.style.SUCCESS(f'成功加载 {len(movies_df)} 部电影')
        )
    
    def load_ratings(self, ratings_file):
        """加载评分数据"""
        self.stdout.write('正在加载评分数据...')
        
        # 读取评分数据
        ratings_df = pd.read_csv(ratings_file)
        
        # 限制加载数量以避免过长时间
        max_ratings = 10000
        if len(ratings_df) > max_ratings:
            self.stdout.write(f'数据量较大，仅加载前 {max_ratings} 条评分数据')
            ratings_df = ratings_df.head(max_ratings)
        
        with transaction.atomic():
            for _, row in ratings_df.iterrows():
                user_id = row['userId']
                movie_id = row['movieId']
                rating = row['rating']
                
                # 创建或获取用户
                username = f'user_{user_id}'
                user, created = User.objects.get_or_create(
                    username=username,
                    defaults={
                        'email': f'{username}@example.com',
                        'first_name': f'User',
                        'last_name': str(user_id)
                    }
                )
                
                # 获取电影
                try:
                    movie = Movie.objects.get(movie_id=movie_id)
                except Movie.DoesNotExist:
                    self.stdout.write(
                        self.style.WARNING(f'电影 {movie_id} 不存在，跳过评分')
                    )
                    continue
                
                # 创建评分
                Rating.objects.get_or_create(
                    user=user,
                    movie=movie,
                    defaults={'rating': rating}
                )
        
        self.stdout.write(
            self.style.SUCCESS(f'成功加载 {len(ratings_df)} 条评分数据')
        )